import pandas as pd
import numpy as np
import sqlite3
from sklearn.model_selection import train_test_split
from sklearn.linear_model import Ridge
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor, BaggingRegressor
from sklearn.tree import DecisionTreeRegressor
from sklearn.metrics import mean_squared_error


def load_data():
    return pd.read_csv("./dataset.csv")


# 数据处理
data = load_data()
df = data[["居室数", "厅堂数", "卫生间数", "总面积", "建造年份", "居民楼总层数", "小区户数", "小区绿化率", "物业费用",
           "价格"]]
X = data[["居室数", "厅堂数", "卫生间数", "总面积", "建造年份", "居民楼总层数", "小区户数", "小区绿化率", "物业费用"]]
# print(X)
y = data["价格"]

# 缺失值处理
print(df.isnull().sum())
print(X.isnull().sum())
print(df.eq(0).sum())
print(y.eq(0).sum())

# # 创建数据库, 创建数据表并插入数据
conn = sqlite3.connect("house.db")
cur = conn.cursor()
cur.execute("""
CREATE TABLE IF NOT EXISTS house_data (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    bedroom_num integer,
    living_num integer,
    bathroom_num integer,
    total_area real,
    build_year integer,
    total_floors integer,
    house_num integer,
    green real,
    manage_fee real,
    price REAL
)
""")
conn.commit()

# df.to_sql("house_data", conn, if_exists="append", index=False)

# for index, row in df.iterrows():
#     params = list(row.values)
#     # print(params)
#     cur.execute("""
#     INSERT INTO house_data (bedroom_num, living_num, bathroom_num, total_area, build_year,
#     total_floors, house_num, green, manage_fee, price) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
#     """, params)
# conn.commit()

# 标准化
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# print(X_scaled)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建模型并训练
estimator1 = LinearRegression()
estimator1.fit(X_train, y_train)
y_pred1 = estimator1.predict(X_test)
estimator2 = RandomForestRegressor(n_estimators=100, random_state=42)
estimator2.fit(X_train, y_train)
y_pred2 = estimator2.predict(X_test)
estimator3 = BaggingRegressor(estimator=DecisionTreeRegressor(), n_estimators=50, random_state=42)
estimator3.fit(X_train, y_train)
y_pred3 = estimator3.predict(X_test)


# 评估模型
mse1 = mean_squared_error(y_test, y_pred1)
mse2 = mean_squared_error(y_test, y_pred2)
print("Linear Regression RMSE:", np.sqrt(mse1))
print("Random Forest RMSE:", np.sqrt(mse2))
mse3 = mean_squared_error(y_test, y_pred3)
print("Bagging Regression RMSE:", np.sqrt(mse3))



# 模型保存
import joblib

joblib.dump(estimator1, "linear_regression.pkl")
joblib.dump(estimator2, "random_forest.pkl")
joblib.dump(estimator3, "bagging_regression.pkl")

